{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "48ef500a-3745-4c78-b848-cccdd110cf13",
   "metadata": {},
   "source": [
    "# Small-to-big Retrieval Pack\n",
    "\n",
    "This LlamaPack provides an example of our small-to-big retrieval (with recursive retrieval)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "006dda50-623c-4292-8e97-f21009405296",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f98d8a24-e225-42a1-9fe8-cbccc776d345",
   "metadata": {},
   "source": [
    "### Setup Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caca86f9-03da-4e51-a308-a2501476a53a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget \"https://www.dropbox.com/s/f6bmb19xdg0xedm/paul_graham_essay.txt?dl=1\" -O paul_graham_essay.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f77671a7-34a0-4263-949b-3bc60b8e219c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import SimpleDirectoryReader\n",
    "from llama_index.node_parser import SimpleNodeParser\n",
    "\n",
    "# load in some sample data\n",
    "reader = SimpleDirectoryReader(input_files=[\"paul_graham_essay.txt\"])\n",
    "documents = reader.load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee30ad51-f3a2-4563-8490-73660175161a",
   "metadata": {},
   "source": [
    "### Download and Initialize Pack\n",
    "\n",
    "Note that this pack directly takes in the html file, no need to load it beforehand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4ebd4574-f56c-4c3b-9bce-49a866d254ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llama_pack import download_llama_pack\n",
    "\n",
    "RecursiveRetrieverSmallToBigPack = download_llama_pack(\n",
    "    \"RecursiveRetrieverSmallToBigPack\",\n",
    "    \"./recursive_retriever_stb_pack\",\n",
    "    # leave the below commented out (was for testing purposes)\n",
    "    # llama_hub_url=\"https://raw.githubusercontent.com/run-llama/llama-hub/jerry/add_llama_packs/llama_hub\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "261eb4d3-726d-483a-b03b-ef65657fd122",
   "metadata": {},
   "outputs": [],
   "source": [
    "recursive_retriever_stb_pack = RecursiveRetrieverSmallToBigPack(\n",
    "    documents,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bb88c53-a0aa-4dea-8ac9-4d6e77fff9d5",
   "metadata": {},
   "source": [
    "### Run Pack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0b69bf0a-6f71-493b-a9aa-3ad0a1e9e0e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;34mRetrieving with query id None: What did the author do growing up?\n",
      "\u001b[0m\u001b[1;3;38;5;200mRetrieved node with id, entering: node-0\n",
      "\u001b[0m\u001b[1;3;34mRetrieving with query id node-0: What did the author do growing up?\n",
      "\u001b[0m"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "source": [
    "# this will run the full pack\n",
    "response = recursive_retriever_stb_pack.run(\"What did the author do growing up?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "86a2eeba-0865-4776-b8b8-a789b8f33435",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author wrote short stories and also worked on programming, specifically on an IBM 1401 computer in 9th grade. They used an early version of Fortran and had to type programs on punch cards. They also mentioned getting a microcomputer, a TRS-80, in about 1980 and started programming on it.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b5e67881-c904-42eb-9ce4-53c5a4d8736d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(response.source_nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd8eef07-df26-4269-aae3-11bdabed704d",
   "metadata": {},
   "source": [
    "### Inspect Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e4e6f2cf-5a87-48b6-9301-d59852430ed1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'query_engine': <llama_index.query_engine.retriever_query_engine.RetrieverQueryEngine at 0x31d055900>,\n",
       " 'recursive_retriever': <llama_index.retrievers.recursive_retriever.RecursiveRetriever at 0x31d056a70>,\n",
       " 'llm': OpenAI(callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x30fd7c0d0>, model='gpt-3.5-turbo', temperature=0.1, max_tokens=None, additional_kwargs={}, max_retries=3, timeout=60.0, default_headers=None, api_key='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', api_base='https://api.openai.com/v1', api_version=''),\n",
       " 'embed_model': HuggingFaceEmbedding(model_name='BAAI/bge-small-en', embed_batch_size=10, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x30fd7c0d0>, tokenizer_name='BAAI/bge-small-en', max_length=512, pooling=<Pooling.CLS: 'cls'>, normalize=True, query_instruction=None, text_instruction=None, cache_folder=None),\n",
       " 'service_context': ServiceContext(llm_predictor=LLMPredictor(system_prompt=None, query_wrapper_prompt=None, pydantic_program_mode=<PydanticProgramMode.DEFAULT: 'default'>), prompt_helper=PromptHelper(context_window=4096, num_output=256, chunk_overlap_ratio=0.1, chunk_size_limit=None, separator=' '), embed_model=HuggingFaceEmbedding(model_name='BAAI/bge-small-en', embed_batch_size=10, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x30fd7c0d0>, tokenizer_name='BAAI/bge-small-en', max_length=512, pooling=<Pooling.CLS: 'cls'>, normalize=True, query_instruction=None, text_instruction=None, cache_folder=None), transformations=[SentenceSplitter(include_metadata=True, include_prev_next_rel=True, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x30fd7c0d0>, chunk_size=1024, chunk_overlap=200, separator=' ', paragraph_separator='\\n\\n\\n', secondary_chunking_regex='[^,.;。？！]+[,.;。？！]?')], llama_logger=<llama_index.logger.base.LlamaLogger object at 0x30fc7f640>, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x30fd7c0d0>)}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "modules = recursive_retriever_stb_pack.get_modules()\n",
    "display(modules)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9cece1d-ae36-4c5c-ab49-6efacd9d761b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_hub",
   "language": "python",
   "name": "llama_hub"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
